
#include <fstream>
#include <iomanip>
#include <sstream>

/*! hypre */
#include "HYPRE.h"
#include "HYPRE_krylov.h"
#include "HYPRE_parcsr_ls.h"
#include "HYPRE_parcsr_mv.h"

/*! unap */
#include "csrMatrix.hpp"
#include "utilities.h"

/*! 定位文件绝对路径相关宏 */
#define LOCATEMATRIX(newName, pre, suff, dir)                         \
  {                                                                   \
    std::ostringstream os;                                            \
    os << dir << pre << (pre.empty() ? "" : "_") << "matrix_" << suff \
       << ".dat";                                                     \
    strcpy(newName, os.str().c_str());                                \
  }
#define LOCATEX(newName, pre, suff, dir)                                   \
  {                                                                        \
    std::ostringstream os;                                                 \
    os << dir << pre << (pre.empty() ? "" : "_") << "vector_init_" << suff \
       << ".dat";                                                          \
    strcpy(newName, os.str().c_str());                                     \
  }
#define LOCATEB(newName, pre, suff, dir)                                  \
  {                                                                       \
    std::ostringstream os;                                                \
    os << dir << pre << (pre.empty() ? "" : "_") << "vector_rhs_" << suff \
       << ".dat";                                                         \
    strcpy(newName, os.str().c_str());                                    \
  }
/*! * * * * * * * * * * * * * * *  * * * * * * * * * * * * * */
using namespace UNAP;
/*! 用于排序的辅助函数 */
template <typename Cmpt>
inline static bool compareTri(const Triplet<Cmpt>& a, const Triplet<Cmpt>& b) {
  if (a.r_ != b.r_) return (a.r_ < b.r_);
  return (a.c_ < b.c_);
};

/*! 确定文件行数 */
int getLineNumber(const char* filename) {
  FILE* fp = fopen(filename, "r");
  label line = 0;
  const int maxCline = 256;
  char cline[maxCline];
  while (fgets(cline, maxCline, fp) != NULL) line++;
  fclose(fp);
  return line;
}
/*! 确定矩阵按进程划分 */
void getMyrange(label& myBeginInd, label& myEndInd, label tsize, int myId,
                int nProcs) {
  label avgColNum = tsize / nProcs;
  label remainder = tsize % nProcs;
  if (myId < remainder) {
    myBeginInd = myId * (avgColNum + 1);
    myEndInd = (myId + 1) * (avgColNum + 1);
  } else {
    myBeginInd = myId * avgColNum + remainder;
    myEndInd = (myId + 1) * avgColNum + remainder;
  }
}
/*! 读取文件构建向量 */
void buildVecFromFile(scalarVector& vec, label myBeginInd, label myEndInd,
                      const char* filename) {
  FILE* fp = fopen(filename, "r");
  if (fp == NULL) {
    std::cerr << "ERROR: could not read file";
    exit(1);
  }
  const int max_cline = 256;
  char cline[max_cline];
  if (fgets(cline, max_cline, fp) == NULL) { /*!< skip first line */
    std::cerr << "ERROR: could not read from file" << std::endl;
    exit(1);
  }
  scalar* data = vec.data();
  scalar bData;
  label count = 0;
  while (fscanf(fp, "%lf\n", &bData) != EOF) {
    if (count >= myBeginInd && count < myEndInd) {
      data[count - myBeginInd] = bData;
    }
    count++;
    if (count >= myEndInd) break;
  }
  fclose(fp);
}

void buildVecFromFile2(scalarVector& vec, label myBeginInd, label myEndInd,
                      const char* filename) {
  FILE* fp = fopen(filename, "r");
  if (fp == NULL) {
    std::cerr << "ERROR: could not read file";
    exit(1);
  }
  const int max_cline = 256;
  char cline[max_cline];
  if (fgets(cline, max_cline, fp) == NULL) { /*!< skip first line */
    std::cerr << "ERROR: could not read from file" << std::endl;
    exit(1);
  }
  scalar* data = vec.data();
  scalar bData;
  label ind,count = 0;
  while (fscanf(fp, "%d %lf\n", &ind, &bData) != EOF) {
    if (count >= myBeginInd && count < myEndInd) {
      data[count - myBeginInd] = bData;
    }
    count++;
    if (count >= myEndInd) break;
  }
  fclose(fp);
}

/*! 读取文件构建矩阵 */
void buildMatFromFile(CSRMatrix<scalar>& A, label myBeginInd, label myEndInd,
                      const char* filename) {
  const int myId = A.comm()->getMyId();
  const int nProcs = A.comm()->getMySize();
  FILE* fp = fopen(filename, "r");
  if (fp == NULL) {
    std::cerr << "ERROR: could not read file";
    exit(1);
  }
  A.setSymm(false);

  std::vector<Triplet<scalar> > Atuple;
  Atuple.reserve(A.nnz() / nProcs);
  const bool zeroBased = 0;
  label ir, ic;
  scalar dv;
  /*!
   * 为了符合mtx矩阵格式，在矩阵最开头加了两行，若矩阵没有这两行则删去下面的skip
   */
  fscanf(fp, "%*[^\n]\n"); /*! skip line */
  fscanf(fp, "%*[^\n]\n"); /*! skip line */
  while (fscanf(fp, "%d %d %lf\n", &ir, &ic, &dv) != EOF) {
    if (ir >= myBeginInd + !zeroBased && ir < myEndInd + !zeroBased) {
      Atuple.push_back(Triplet<scalar>(ir, ic, dv));
    }
  }
  fclose(fp);

  A.lRows() = A.lCols() = myEndInd - myBeginInd;
  A.lnnz() = Atuple.size();
  if (!zeroBased)
    for (label i = 0; i < A.lnnz(); ++i) {
      (Atuple[i].r_)--;
      (Atuple[i].c_)--;
    }

  A.beginRow() = A.beginCol() = myBeginInd;
  A.vtxdist().resize(nProcs + 1);
  label avgColNum = A.rows() / nProcs;
  label remainder = A.rows() % nProcs;
  for (label i = 0; i < nProcs; ++i) {
    if (i < remainder) {
      A.vtxdist(i) = i * (avgColNum + 1);
    } else {
      A.vtxdist(i) = i * avgColNum + remainder;
    }
  }
  A.vtxdist(nProcs) = A.rows();

  std::stable_sort(Atuple.begin(), Atuple.end(), compareTri<scalar>);
  A.ptr() = new label[A.lRows() + 1];
  A.ind() = new label[A.lnnz()];
  A.val() = new scalar[A.lnnz()];
  label row = -1, count = 0;
  for (label i = 0; i < A.lnnz(); i++) {
    A.val(i) = Atuple[i].v_;
    A.ind(i) = Atuple[i].c_;
    label newRow = Atuple[i].r_;
    if (newRow != row) {
      A.ptr(count) = i;
      row = newRow;
      count++;
    }
  }
  A.ptr(A.lRows()) = A.lnnz();
  A.splitOffDiag();
}
/*! 向量输出工具 */
#define FILEOPEN(fout, fileName, MYID)     \
  std::ostringstream os;                   \
  os << fileName << "_" << MYID << ".txt"; \
  fout.open(os.str().c_str())
#define FILECLOSE(fout, fileName) fout.close()
template <typename T>
void printVector(const T& v, const char* fileName) {
  const label size = v.size();
  std::ofstream fout;

  FILEOPEN(fout, fileName, v.getCommunicator()->getMyId());

  fout << "rowSize: " << size << std::endl;
  forAll(i, size) {
    fout << std::setiosflags(std::ios::scientific) << std::setprecision(15)
         << v[i] << std::endl;
  }
  FILECLOSE(fout, fileName);
}

/*! * * * * * * * * * * * * * * * 主函数  * * * * * * * * * * * * * */
int main(int argc, char* argv[]) {
  /* Initialize MPI */
  COMM::createLogFile(false);
  initUtility();
  Communicator* comm = &COMM::getGlobalComm();
  const int nProcs = comm->getMySize();
  const int myId = comm->getMyId();
  label gSize, size, gNnz, nnz, beginR, endR;

  char dir[200];
  char fileName[200];
  string field;
  if (argc > 1)
    strcpy(dir, argv[1]);
  else
    strcpy(dir, "../kvlcc2/");
  if (argc > 2)
    field = argv[2];
  else
    field = "u";

  string pre;
  if (argc > 3)
    pre = argv[3];
  else
    pre = "";

  string suff = field + "_101";
  /*!< get basic info */
  LOCATEB(fileName, pre, suff, dir)
  gSize = getLineNumber(fileName) - 2;
  getMyrange(beginR, endR, gSize, myId, nProcs);
  size = endR - beginR;
  scalarVector xx(size, comm), bb(size, comm);

  /*!< read vector */
  buildVecFromFile(bb, beginR, endR, fileName);
  LOCATEX(fileName, pre, suff, dir)
  buildVecFromFile2(xx, beginR, endR, fileName);

  /*! read matrix */
  LOCATEMATRIX(fileName, pre, suff, dir);
  gNnz = getLineNumber(fileName);
  CSRMatrix<scalar> A_CSR(comm);
  A_CSR.rows() = gSize;
  A_CSR.nnz() = gNnz;
  buildMatFromFile(A_CSR, beginR, endR, fileName);

  /*! * * * * * * * * * * * * * * * 开始计算 * * * * * * * * * * * * * */
  comm->barrier();
  double start_CPU, end_CPU;
  /*! * * * * * * * * * * * * * * * hypre  * * * * * * * * * * * * * */
  HYPRE_IJMatrix A;
  HYPRE_ParCSRMatrix parcsr_A;
  HYPRE_IJVector b;
  HYPRE_ParVector par_b;
  HYPRE_IJVector x;
  HYPRE_ParVector par_x;
  HYPRE_Solver solver, precond;
  HYPRE_Init();

  if (nProcs > 1) {
    comm->barrier();
  }
  /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
  /*  Create the Vecs and Mat                                            */
  /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
  label Istart = A_CSR.beginRow();
  label Iend = A_CSR.endRow();
  label* rows = A_CSR.ptr();
  label* coo_j = A_CSR.ind();
  scalar* coo_v = A_CSR.val();

  /* Create the matrix.
   | Note that this is a square matrix, so we indicate the row partition
   | size twice (since number of rows = number of cols) */
  HYPRE_IJMatrixCreate(MPI_COMM_WORLD, Istart, Iend - 1, Istart, Iend - 1, &A);
  /* Choose a parallel csr format storage (see the User's Manual) */
  HYPRE_IJMatrixSetObjectType(A, HYPRE_PARCSR);
  /* Initialize before setting coefficients */
  HYPRE_IJMatrixInitialize(A);
  /* set the values for matrix */
  const label rowSize = A_CSR.lRows();
  std::vector<label> nnzPerRow(rowSize), rowId(rowSize);
  forAll(i, rowSize) {
    nnzPerRow[i] = rows[i + 1] - rows[i];
    rowId[i] = i + Istart;
  }

  HYPRE_IJMatrixSetValues(A, Iend - Istart, nnzPerRow.data(), rowId.data(),
                          coo_j, coo_v);

  /* Assemble after setting the coefficients */
  HYPRE_IJMatrixAssemble(A);
  /* Get the parcsr matrix object to use */
  HYPRE_IJMatrixGetObject(A, (void**)&parcsr_A);
  /* Create the rhs and solution */
  HYPRE_IJVectorCreate(MPI_COMM_WORLD, Istart, Iend - 1, &b);
  HYPRE_IJVectorSetObjectType(b, HYPRE_PARCSR);
  HYPRE_IJVectorInitialize(b);
  HYPRE_IJVectorCreate(MPI_COMM_WORLD, Istart, Iend - 1, &x);
  HYPRE_IJVectorSetObjectType(x, HYPRE_PARCSR);
  HYPRE_IJVectorInitialize(x);
  {
    label* rows;

    rows = (label*)calloc(Iend - Istart, sizeof(label));

    for (label i = 0; i < Iend - Istart; i++) {
      rows[i] = Istart + i;
    }

    HYPRE_IJVectorSetValues(b, Iend - Istart, rows, bb.data());
    HYPRE_IJVectorSetValues(x, Iend - Istart, rows, xx.data());

    free(rows);
  }
  HYPRE_IJVectorAssemble(b);
  HYPRE_IJVectorGetObject(b, (void**)&par_b);
  HYPRE_IJVectorAssemble(x);
  HYPRE_IJVectorGetObject(x, (void**)&par_x);
  // HYPRE_IJMatrixPrint(A, "IJ.out.A");
  //  HYPRE_IJVectorPrint(b, "IJ.out.b");

  /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
  /*  Solve the linear system of equations                               */
  /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */

  start_CPU = MPI_Wtime();  // 计时函数
  label num_iterations;
  scalar final_res_norm;
  /* Create solver */
  HYPRE_BoomerAMGCreate(&solver);
  HYPRE_BoomerAMGSetTol(solver, 1.e-6); /* conv. tolerance */
  // HYPRE_BoomerAMGSetMaxCoarseSize(solver, 50);
  HYPRE_BoomerAMGSetMaxIter(solver, 100);
  // HYPRE_BoomerAMGSetMaxLevels(solver, 50); /* maximum number of levels */
  /* Set some parameters (See Reference Manual for more parameters) */
  HYPRE_BoomerAMGSetPrintLevel(solver, 3); /* print solve info + parameters */
  //HYPRE_BoomerAMGSetRelaxType(solver, 0);  /// jacobi
  // HYPRE_BoomerAMGSetRelaxOrder(solver, false);
  //HYPRE_BoomerAMGSetCoarsenType(solver, 10);  /// PMIS-8  HMIS-10
  // HYPRE_BoomerAMGSetInterpType(solver, 3); /* 3, 15, 6, 14, 18 */
  // HYPRE_BoomerAMGSetAggNumLevels(solver, agg_num_levels);
  // HYPRE_BoomerAMGSetAggInterpType(solver, 5); /* 5 or 7 */
  // HYPRE_BoomerAMGSetKeepTranspose(solver, true); /* keep transpose to avoid
  //  SpMTV */
  /* Now setup and solve! */

  /*! setup  */
  HYPRE_BoomerAMGSetup(solver, parcsr_A, par_b, par_x);

  /*! solve */
  HYPRE_BoomerAMGSolve(solver, parcsr_A, par_b, par_x);

  /* Run info - needed logging turned on */
  HYPRE_BoomerAMGGetNumIterations(solver, &num_iterations);
  HYPRE_BoomerAMGGetFinalRelativeResidualNorm(solver, &final_res_norm);

  HYPRE_IJVectorPrint(x, field.c_str());
  /* Destroy solver */
  HYPRE_BoomerAMGDestroy(solver);

  /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
  /*  Free up memory                                                     */
  /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
  /* Clean up */
  HYPRE_IJMatrixDestroy(A);
  HYPRE_IJVectorDestroy(b);
  HYPRE_IJVectorDestroy(x);
  /* Finalize HYPRE */
  HYPRE_Finalize();
  /* Finalize MPI*/

  /*! * * * * * * * * * * * * * * *  * * * * * * * * * * * * * */

  end_CPU = MPI_Wtime();  // 计时函数
  double time_CPU = ((double)end_CPU - (double)start_CPU);
  MPI_Allreduce(MPI_IN_PLACE, &time_CPU, 1, MPI_DOUBLE, MPI_MAX,
                comm->getMpiComm());

  COUT << "\nAfter " << num_iterations
       << " iterations, the solution is converged!" << ENDL
       << "final relative Res " << final_res_norm << ENDL;

  if (myId == 0) {
    printf("max_time_CPU=%f\n", time_CPU);
  }
}
